package map;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Objects;
import java.util.Queue;

public class TreeMap<K, V> implements Map<K, V>{
	
	private static final boolean RED = false;
	private static final boolean BLACK = true;

	private int size;
	private Node<K, V> root;
	private Comparator<K> comparator;
	
	@Override
	public int size() {
		return size;
	}

	@Override
	public boolean isEmpty() {
		return size == 0;
	}

	@Override
	public void clear() {
		size = 0;
		root = null;
	}

	@Override
	public V put(K key, V value) {
		keyNotNullCheck(key);
		if (root == null) {
			root = new Node<>(key, value, null);
			size++;
			afterPut(root);
			return null;
		}
		Node<K, V> node = root;
		Node<K, V> parent = root;
		int cmp = 0;
		while (node != null) {
			parent = node;
			cmp = compare(key, node.key);
			if (cmp > 0) {
				node = node.right;
			} else if (cmp < 0) {
				node = node.left;
			} else {
				node.key = key;
				V oldValue = value;
				node.value = value;
				return oldValue;
			}
		}
		Node<K, V> newNode = new Node<>(key, value, parent);
		if (cmp > 0) {
			parent.right = newNode;
		} else {
			parent.left = newNode;
		}
		size++;
		afterPut(newNode);
		return null;
	}
	
	private void afterPut(Node<K, V> node) {
		Node<K, V> parent = node.parent;
		if (parent == null) {
			black(node);
			return;
		}
		if (isBlack(parent)) return;
		
		Node<K, V> uncle = parent.sibling();
		Node<K, V> grand = red(parent.parent);
		if (isRed(uncle)) {
			black(parent);
			black(uncle);
			afterPut(grand);
			return;
		} 
		
		if (parent.isLeftChild()) {
			if (node.isRightChild()) {
				black(node);
				rotateLeft(parent);
			} else {
				black(parent);
			}
			rotateRight(grand);
		} else {
			if (node.isLeftChild()) {
				black(node);
				rotateRight(parent);
			} else {
				black(parent);
			}
			rotateLeft(grand);
		}
	}
	
	
	@Override
	public V get(K key) {
		Node<K, V> node = node(key);
		return node == null ? null : node.value;
	}

	@Override
	public V remove(K key) {
		return remove(node(key));
	}
	
	private V remove(Node<K, V> node) {
		if (node == null) return null;
		size--;
		
		V oldValue = node.value;
		
		if (node.hasTwoChildren()) {
			Node<K, V> sucNode = successor(node);
			node.key = sucNode.key;
			node.value = sucNode.value;
			node = sucNode;
		}
		
		Node<K, V> replacement = node.left != null ? node.left : node.right;
		
		if (replacement != null) {
			replacement.parent = node.parent;
			if (node.parent == null) {
				root = replacement;
			} else if (node == node.parent.left) {
				node.parent.left = replacement;
			} else {
				node.parent.right = replacement;
			}
			afterRemove(replacement);
		} else if (node.parent == null) {
			root = null;
		} else {
			if (node == node.parent.left) {
				node.parent.left = null;
			} else {
				node.parent.right = null;
			}
			afterRemove(node);
		} 
		return oldValue;
	}
	
	protected void afterRemove(Node<K, V> node) {
		if (isRed(node)) {
			black(node);
			return;
		}
		Node<K, V> parent = node.parent;
		if (parent == null) return;
		boolean left = parent.left == null || node.isLeftChild();
		Node<K, V> sibling = left ? parent.right : parent.left;
		if (left) {
			if (isRed(sibling)) {
				black(sibling);
				red(parent);
				rotateLeft(parent);
				sibling = parent.right;
			}
			if (isBlack(sibling.left) && isBlack(sibling.right)) {
				boolean parentBlack = isBlack(parent);
				black(parent);
				red(sibling);
				if (parentBlack) {
					afterRemove(parent);
				}
			} else {
				if (isBlack(sibling.right)) {
					rotateRight(sibling);
					sibling = parent.right;
				}
				color(sibling, colorOf(parent));
				black(sibling.right);
				black(parent);
				rotateLeft(parent);
			}
			
		} else {
			if (isRed(sibling)) {
				black(sibling);
				red(parent);
				rotateRight(parent);
				sibling = parent.left;
			}
			if (isBlack(sibling.left) && isBlack(sibling.right)) {
				boolean parentBlack = isBlack(parent);
				black(parent);
				red(sibling);
				if (parentBlack) {
					afterRemove(parent);
				}
			} else {
				if (isBlack(sibling.left)) {
					rotateLeft(sibling);
					sibling = parent.left;
				}
				color(sibling, colorOf(parent));
				black(sibling.left);
				black(parent);
				rotateRight(parent);
			}
			
		}
	}
	
	@Override
	public boolean containsKey(K key) {
		return node(key) != null;
	}

	@Override
	public boolean containsValue(V value) {
		if (root == null) return false;
		Queue<Node<K, V>> queue = new LinkedList<>();
		queue.offer(root);
		while (!queue.isEmpty()) {
			Node<K, V> node = queue.poll();
			if (Objects.equals(value, node.value)) return true;
			if (node.left != null) {
				queue.offer(node.left);
			}
			
			if (node.right != null) {
				queue.offer(node.right);
			}
		}
		return false;
	}

	@Override
	public void traversal(Visitor<K, V> visitor) {
		if (visitor == null) return;
		traversal(visitor, root);
	}
	
	private void traversal(Visitor<K, V> visitor, Node<K, V> node) {
		if (node == null || visitor.stop) return;
		traversal(visitor, node.left);
		if (visitor.stop) return;
		visitor.stop = visitor.visit(node.key, node.value);
		traversal(visitor, node.right);
	}
	
	private Node<K, V> successor(Node<K, V> node) {
		if (node == null) return null;
		Node<K, V> sucNode = node.right;
		if (sucNode != null) {
			while (sucNode.left != null) {
				sucNode = sucNode.left;
			}
			return sucNode;
		}
		while (node.parent != null && node == node.parent.right) {
			node = node.parent;
		}
		return node.parent;
	}
	
	private Node<K, V> node(K key) {
		keyNotNullCheck(key);
		Node<K, V> node = root;
		int cmp = 0;
		while (node != null) {
			cmp = compare(key, node.key);
			if (cmp == 0) {
				return node;
			} else if (cmp > 0) {
				node = node.right;
			} else {
				node = node.left;
			}
		}
		return null;
	}
	
	private int compare(K k1, K k2) {
		if (comparator != null) return comparator.compare(k1, k2);
		return ((Comparable<K>)k1).compareTo(k2);
	}
	
	private void keyNotNullCheck(K key) {
		if (key == null) throw new NullPointerException("Key must not be null");
	}
	
	private Node<K, V> red(Node<K, V> node) {
		return color(node, RED);
	}
	
	private Node<K, V> black(Node<K, V> node) {
		return color(node, BLACK);
	}
	
	private Node<K, V> color(Node<K, V> node, boolean color) {
		if (node == null) return null;
		((Node<K, V>)node).color = color;
		return node;
	}
	
	private boolean isBlack(Node<K, V> node) {
		return colorOf(node) == BLACK;
	}
	
	private boolean isRed(Node<K, V> node) {
		return colorOf(node) == RED;
	}
	
	private boolean colorOf(Node<K, V> node) {
		return node == null ? BLACK : ((Node<K, V>)node).color;
	}
	
	protected void rotateLeft(Node<K, V> grand) {
		Node<K, V> parent = grand.right;
		Node<K, V> child = parent.left;
		grand.right = child;
		parent.left = grand;
		afterRotate(grand, parent, child);
	}

	private void rotateRight(Node<K, V> grand) {
		Node<K, V> parent = grand.left;
		Node<K, V> child = parent.right;
		grand.left = child;
		parent.right = grand;
		afterRotate(grand, parent, child);
	}

	private void afterRotate(Node<K, V> grand, Node<K, V> parent, Node<K, V> child) {
		parent.parent = grand.parent;
		if (grand.isLeftChild()) {
			grand.parent.left = parent;
		} else if (grand.isRightChild()) {
			grand.parent.right = parent;
		} else {
			root = parent;
		}
		
		if (child != null) {
			child.parent = grand;
		}
		
		grand.parent = parent;
	}
	
	private static class Node<K, V> {
		K key;
		V value;
		Node<K, V> parent;
		Node<K, V> left;
		Node<K, V> right;
		boolean color = RED;
		
		public Node(K key, V value, Node<K, V> parent) {
			this.key = key;
			this.value = value;
			this.parent = parent;
		}
		
		public boolean isLeaf() {
			return this.left == null && this.right == null;
		}
		
		public boolean hasTwoChildren() {
			return this.left != null && this.right != null;
		}
		
		public boolean isLeftChild() {
			return parent != null && this == parent.left;
		}
		public boolean isRightChild() {
			return parent != null && this == parent.right;
		}
		
		public Node<K, V> sibling() {
			if (isLeftChild()) return parent.right;
			if (isRightChild()) return parent.left;
			return null;
		}
	}

}
